from sklearn.svm import SVC
from hog import *
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
import itertools
from sklearn.metrics import confusion_matrix


def plot_confusion_matrix(cm, classes, title='Confusion Matrix'):
    plt.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)
    plt.title(title)
    tick_marks = np.arange(len(classes))
    plt.xticks(tick_marks, classes, rotation=0)
    plt.yticks(tick_marks, classes)
    thresh = cm.max() / 2.
    # add digits
    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
        plt.text(j, i, cm[i, j],
                 horizontalalignment="center",
                 color="white" if cm[i, j] > thresh else "black")
    plt.tight_layout()
    plt.ylabel('True Label')
    plt.xlabel('Predicted Label')
    plt.show()


def batch_iter(x, y, batch_size, shuffle=True):
    data_size = len(x)
    num_batches = ((data_size-1)//batch_size) + 1
    if shuffle:
        shuffle_indices = np.random.permutation(np.arange(data_size))
        shuffled_x = x[shuffle_indices]
        shuffled_y = y[shuffle_indices]
    else:
        shuffled_x = x
        shuffled_y = y
    for batch_num in range(num_batches):
        start_index = batch_num*batch_size
        end_index = min((batch_num+1)*batch_size, data_size)
        yield shuffled_x[start_index:end_index], shuffled_y[start_index:end_index]


def load_data(path):
    return pickle.load(open(path, 'rb'))


def main():
    features = load_data('data/features.pkl')
    labels = load_data('data/labels.pkl')
    clf = SVC()
    x_train, x_test, y_train, y_test = train_test_split(features, labels, test_size=0.2)
    print('fit _ _ _')
    for e in range(2):
        print(e, '-->')
        for batch_x, batch_y in batch_iter(x_train, y_train, 200):
            clf.fit(batch_x, batch_y)

    pickle.dump(clf, open('data/clf2.pkl', 'wb'))
    pickle.dump(x_test, open('data/x_test.pkl', 'wb'))
    pickle.dump(y_test, open('data/y_test.pkl', 'wb'))
    pass


def predict():
    clf = load_data('data/clf3.pkl')
    x_test = load_data('data/x_test.pkl')
    y_test = load_data('data/y_test.pkl')
    print('predict-->')
    print(y_test.shape)
    y_predict = clf.predict(x_test)
    print('finish predict')
    print('draw cnf')
    cnf_matrix = confusion_matrix(y_test, y_predict)
    print(cnf_matrix)
    print('plot cnf')
    plot_confusion_matrix(cnf_matrix, classes=[0, 1])
    pass


def test_plot():
    predict_demo = np.array([[0], [1], [1], [0], [0], [0], [0], [1], [1], [1]])
    label_demo = np.array([[0], [1], [1], [1], [0], [0], [0], [1], [1], [1]])
    print(predict_demo.shape)
    print(label_demo.shape)
    cnf_matrix = confusion_matrix(label_demo, predict_demo)
    print(cnf_matrix)
    plot_confusion_matrix(cnf_matrix, classes=['dog', 'cat'])
    pass


def predict_img():
    clf = load_data('data/clf3.pkl')
    feature = get_feature('data/1.jpg')
    result = clf.predict([feature])
    if result == 0:
        print('dog')
    if result == 1:
        print('cat')
    pass


def fit_data():
    features = load_data('data/features.pkl')
    labels = load_data('data/labels.pkl')
    clf = SVC()
    x_train, x_test, y_train, y_test = train_test_split(features, labels, test_size=0.2)
    clf.fit(x_train, y_train)
    # print('fit _ _ _')
    for e in range(5):
        print(e, '-->')
        for batch_x, batch_y in batch_iter(x_train, y_train, 200):
            clf.fit(batch_x, batch_y)

    pickle.dump(clf, open('data/clf3.pkl', 'wb'))
    pickle.dump(x_test, open('data/x_test.pkl', 'wb'))
    pickle.dump(y_test, open('data/y_test.pkl', 'wb'))
    pass


if __name__ == '__main__':
    main()
    predict()
    # test_plot()
    # predict_img()
    # fit_data()
